package com.laboratory.sselearning.service;

import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.time.Instant;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class SseService {

    private final Map<String, SseEmitter> clients = new ConcurrentHashMap<>();

    /**
     * 创建SSE连接
     *
     * @param userId
     * @param timeoutMilliseconds
     * @return
     */
    public SseEmitter connect(String userId, Long timeoutMilliseconds) {
        long timeout = timeoutMilliseconds != null ? timeoutMilliseconds : -1L;
        SseEmitter sseEmitter = new SseEmitter(timeout);
        clients.put(userId, sseEmitter);
        /**
         * 监听完成
         */
        sseEmitter.onCompletion(
                () -> {
                    clients.remove(userId);
                    System.out.println("用户 " + userId + " 断开连接");
                }
        );
        /**
         * 监听超时
         */
        sseEmitter.onTimeout(
                () -> {
                    clients.remove(userId);
                    System.out.println("用户 " + userId + " 超时");
                }
        );
        /**
         * 监听错误
         */
        sseEmitter.onError(
                throwable -> {
                    clients.remove(userId);
                    System.out.println("用户 " + userId + " 错误");
                }
        );
        try {
            sseEmitter.send(
                    SseEmitter.event()
                            .id(String.valueOf(System.currentTimeMillis()))
                            .name("INIT")
                            .reconnectTime(3000)
                            .data(Map.of("message", "初始化成功"))
            );
        } catch (Exception e) {
            clients.remove(userId);
            sseEmitter.completeWithError(e);
        }
        return sseEmitter;
    }

    /**
     * 发送消息给指定用户（单播）
     *
     * @param userId
     * @param data
     * @return
     */
    public boolean sendTo(String userId, Object data) {
        SseEmitter sseEmitter = clients.get(userId);
        if (sseEmitter == null) {
            return false;
        }
        try {
            sseEmitter.send(
                    SseEmitter.event()
                            .id(String.valueOf(System.currentTimeMillis()))
                            .name("MESSAGE")
                            .data(data)
            );
            return true;
        } catch (Exception e) {
            clients.remove(userId);
            sseEmitter.completeWithError(e);
            return false;
        }
    }

    /**
     * 广播消息（多播）
     *
     * @param data
     * @return
     */
    public int broadcast(Object data) {
        int count = 0;
        for (var entry : clients.entrySet()) {
            try {
                entry.getValue().send(
                        SseEmitter.event()
                                .name("BROADCAST")
                                .data(data)
                );
                count++;
            } catch (Exception e) {
                clients.remove(entry.getKey());
                entry.getValue().completeWithError(e);
            }
        }
        return count;
    }

    /**
     * 关闭连接
     *
     * @param userId
     */
    public void close(String userId) {
        SseEmitter emitter = clients.remove(userId);
        if (emitter != null) {
            emitter.complete();
        }
    }

    /**
     * 心跳（多播）
     */
    public void beatAll() {
        for (var entry : clients.entrySet()) {
            try {
                entry.getValue().send(
                        SseEmitter.event()
                                .comment("hb@" + Instant.now())
                );
            } catch (Exception e) {
                entry.getValue().completeWithError(e);
                clients.remove(entry.getKey());
            }
        }
    }

}
